
#include "glodef.h"
#include "errlog.h"
#include "grid3d.h"
#include "vecs.h"
#include "fe.h"
#include "fe_matr.h"

#include <rk45.h>


typedef struct
{
    const grid3d_t *grid;
    const fe_matr_t *mass;
    const vector *adv;
    const vector *adv2;
    vector tmp;
    vector tm1;
    unsigned eltype;
    int count;
} rk45_data_t;

static rk45_data_t rk45_data;

/****************************************************************************/
/*       The Runge-Kutta 45 used for the state equation                     */
/****************************************************************************/

static void dydt_state(double *t, double y[], double yp[])
{
    vector u0, u;
    int i;
    (void) t;
    
    rk45_data.count++;
    
    u0.N = rk45_data.adv->N;
    u0.X = y;
    u0.Y = u0.X + u0.N;
    u0.Z = u0.Y + u0.N;
    
    u.N = rk45_data.adv->N;
    u.X = yp;
    u.Y = u.X + u.N;
    u.Z = u.Y + u.N;
    
    test_advection_vector(rk45_data.grid, rk45_data.eltype, rk45_data.adv, &u0, &rk45_data.tmp);
    fe_matr_solve_vector(rk45_data.mass, &u, &rk45_data.tmp, &rk45_data.tm1, 0.0);
    
    for(i=0;i<3*u.N;i++) u.X[i] *= -1;
}

void advection_rk45(const grid3d_t *grid, unsigned eltype, double h, const fe_matr_t *massU,
        const vector *adv_vel, const vector *u0, vector *u)
{
    int iflag, neqn;
    int ndof = fe_ndof(grid, eltype);
    double *work;
    int iwork[5];
    
    clock_t tm = clock();
    
    if(u->N!=ndof || u0->N!=ndof || adv_vel->N!=ndof)
        error_msg("%s: dimensions disagree\n", __func__);
    
    PrintLog("Advection by rk45:  ");
    TurnLogOff();
    
    rk45_data.count = 0;
    rk45_data.adv = adv_vel;
    rk45_data.eltype = eltype;
    rk45_data.grid = grid;
    rk45_data.mass = massU;
    
    alloc_vec(&rk45_data.tm1, u->N);
    alloc_vec(&rk45_data.tmp, u->N);
    
    neqn = 3*ndof;
    work = allocV(3+6*neqn);
    
    copy_vec(u0, u);
    double t = 0.0, tout = h, rerr = 1e-3, aerr = 1e-6;
    iflag = 1;
    while((t < tout) && (iflag < 5))
        rkf45_(dydt_state, &neqn, u->X, &t, &tout, &rerr, &aerr, &iflag, work, iwork);
    
    if(iflag != 2)
        error_msg("%s: rkf45_ failed\n iflag = %d  t = %g  tout = %g\n",
                __func__, iflag, t, tout);
    
    freeV(&work);
    free_vec(&rk45_data.tm1);
    free_vec(&rk45_data.tmp);
    
    TurnLogOn();
    double cpu = (clock()-tm) / (double) CLOCKS_PER_SEC;
    PrintLog("count %d  CPU time=%-4.2fsec.\n", rk45_data.count, cpu);
}



/****************************************************************************/
/*       The Runge-Kutta 45 used for the adjoint equation                     */
/****************************************************************************/

static void dydt_adjoint(double *t, double y[], double yp[])
{
    vector v0, v;
    int i;
    (void) t;
    
    rk45_data.count++;
    
    v0.N = rk45_data.adv->N;
    v0.X = y;
    v0.Y = v0.X + v0.N;
    v0.Z = v0.Y + v0.N;
    
    v.N = rk45_data.adv->N;
    v.X = yp;
    v.Y = v.X + v.N;
    v.Z = v.Y + v.N;
    
    test_advection_vector(rk45_data.grid, rk45_data.eltype, rk45_data.adv, &v0, &rk45_data.tmp);
    test_advection_vector(rk45_data.grid, rk45_data.eltype, &v0, rk45_data.adv2, &rk45_data.tm1);
    for(i=0;i<3*v.N;i++) rk45_data.tmp.X[i] = rk45_data.tm1.X[i] - rk45_data.tmp.X[i];
    fe_matr_solve_vector(rk45_data.mass, &v, &rk45_data.tmp, &rk45_data.tm1, 0.0);
    
    for(i=0;i<3*v.N;i++) v.X[i] *= -1;
}

void advection_rk45_adjoint(const grid3d_t *grid, unsigned eltype, double h, const fe_matr_t *massU,
        const vector *adv_vel, const vector *un, const vector *v1, vector *v)
{
    int iflag, neqn;
    int ndof = fe_ndof(grid, eltype);
    double *work;
    int iwork[5];
    
    clock_t tm = clock();
    
    if(v->N!=ndof || v1->N!=ndof || adv_vel->N!=ndof || un->N!=ndof)
        error_msg("%s: dimensions disagree\n", __func__);
    
    PrintLog("Advection by rk45:  ");
    TurnLogOff();
    
    rk45_data.count = 0;
    rk45_data.adv = adv_vel;
    rk45_data.adv2 = un;
    rk45_data.eltype = eltype;
    rk45_data.grid = grid;
    rk45_data.mass = massU;
    
    alloc_vec(&rk45_data.tm1, v->N);
    alloc_vec(&rk45_data.tmp, v->N);
    
    neqn = 3*ndof;
    work = allocV(3+6*neqn);
    
    copy_vec(v1, v);
    double t = 0.0, tout = h, rerr = 1e-3, aerr = 1e-6;
    iflag = 1;
    while((t < tout) && (iflag < 5))
        rkf45_(dydt_adjoint, &neqn, v->X, &t, &tout, &rerr, &aerr, &iflag, work, iwork);
    
    if(iflag != 2)
        error_msg("%s: rkf45_ failed\n iflag = %d  t = %g  tout = %g\n",
                __func__, iflag, t, tout);
    
    freeV(&work);
    free_vec(&rk45_data.tm1);
    free_vec(&rk45_data.tmp);
    
    TurnLogOn();
    double cpu = (clock()-tm) / (double) CLOCKS_PER_SEC;
    PrintLog("count %d  CPU time=%-4.2fsec.\n", rk45_data.count, cpu);
}


